import os, json, asyncio, websockets, sounddevice as sd
import numpy as np
import serial

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
REALTIME_URL = "wss://api.openai.com/v1/realtime?model=gpt-realtime"  # 或 Azure 同等端点
SR = 16000
CHUNK = 320  # 20ms @16kHz

# === 串口初始化（沿用你的）===
uart = serial.Serial("COM37", 115200, timeout=0)

# === 工具定义：只允许输出坐标/G指令 ===
TOOL_SPEC = {
  "type": "function",
  "name": "emit_move_batch",
  "description": "发送一批机械臂G命令（只包含G0与 X/Y/Z/F 整数或固定F=50）。不得输出其它文本。",
  "parameters": {
    "type": "object",
    "properties": {
      "lines": {
        "type": "array",
        "items": {
          "type": "string",
          "pattern": r"^(G0)\s+X\d{1,3}\s+Y\d{1,3}\s+Z\d{1,3}\s+F(50|\d{2,4})$"
        },
        "minItems": 1, "maxItems": 1024
      }
    },
    "required": ["lines"],
    "additionalProperties": False
  }
}

SYSTEM_PROMPT = (
    "你是机械臂操作员。仅在需要运动时调用工具 emit_move_batch，输出严格的G命令，"
    "不得夹带解释性文字。X/Y/Z取值0-250的整数，F固定为50除非我另行指明。"
    "例如：G0 X10 Y160 Z150 F50 \n"
)

async def audio_producer(ws):
    """实时从麦克风采样并发送到 Realtime 会话"""
    def callback(indata, frames, time, status):
        if status:  # 采样状态警告
            return
        # 16kHz 16-bit PCM 小块
        pcm16 = (indata[:, 0] * 32767).astype(np.int16).tobytes()
        # 发送二进制音频帧（Realtime 接口支持 audio chunks）
        asyncio.get_event_loop().create_task(ws.send(pcm16))

    with sd.InputStream(samplerate=SR, channels=1, dtype='float32',
                        blocksize=CHUNK, callback=callback):
        while True:
            await asyncio.sleep(0.02)  # 与块时长一致

async def message_consumer(ws):
    """消费模型事件：当出现工具调用就落地到串口"""
    async for msg in ws:
        if isinstance(msg, bytes):
            # 可能是音频输出，与你无关，可忽略
            continue
        data = json.loads(msg)
        # Realtime 会把工具调用作为一个事件（名称和具体字段以文档为准）
        if data.get("type") == "tool_call" and data.get("name") == "emit_move_batch":
            args = data.get("arguments", {})
            lines = args.get("lines", [])
            for line in lines:
                uart.write((line + "\n").encode("utf-8"))

async def main():
    headers = [
        ("Authorization", f"Bearer {OPENAI_API_KEY}"),
        ("OpenAI-Beta", "realtime=v1")  # 视文档要求
    ]
    async with websockets.connect(REALTIME_URL, extra_headers=headers, ping_interval=20) as ws:
        # 1) 发送系统提示与工具声明
        await ws.send(json.dumps({
            "type": "session.update",
            "instructions": SYSTEM_PROMPT,
            "tools": [TOOL_SPEC],
            # 可选：开启模型端端点检测/说话人中断等加速
            "turn_detection": {"type": "server_vad", "silence_ms": 300}
        }))
        # 2) 并发跑音频生产与消息消费
        await asyncio.gather(audio_producer(ws), message_consumer(ws))

if __name__ == "__main__":
    asyncio.run(main())
